/*
 * Copyright (c) 2025.
 * All rights reserved.
 */

#include "driver.h"
#include "path_filter.h"
#include "string_utils.h"

// size of '\?' is 2
#define PATH_START_OFFSET       2
// size of '\??\C:\' is 7
#define PATH_FOLDER_OFFSET      7
// size of '\??\C:\Progra' is 13
#define PATH_PROGRA_OFFSET      13

#define PATH_MAX_SIZE_SHORT     14
// size of '\??\C:\Progra~1\a.exe' is 21
// size of '\??\C:\Windows\a.exe' is 20,path->Length must >19
#define PATH_MIN_SIZE_EXE       19

// for performance,"Progra"
#define SIZE_OF_PROGRA          6
// "m Files"
#define SIZE_OF_MFILES          7
// "Windows\\"
#define SIZE_OF_WINDOWS         8
// "?\\C:"
#define SIZE_OF_C               4

/// <summary>
/// Compare [\??\C:\] Windows\*
/// </summary>
/// <param name="path"></param>
/// <returns></returns>
BOOLEAN IsInWindows(
    PCUNICODE_STRING path)
{
    if (IsNullOrEmpty(path))
        return FALSE;
    if (_wcsnicmp(path->Buffer + PATH_FOLDER_OFFSET, L"Windows\\", SIZE_OF_WINDOWS) == 0) { //ignore case
        return TRUE;
    }
    return FALSE;
}

/// <summary>
/// Check whether it is a short path
/// </summary>
/// <param name="path"></param>
/// <returns></returns>
BOOLEAN IsShortPath(
    PCUNICODE_STRING path)
{
    if (IsNullOrEmpty(path) || path->Length < PATH_MAX_SIZE_SHORT)
        return FALSE;
    int i = PATH_FOLDER_OFFSET + 1;
    for (; i < PATH_MAX_SIZE_SHORT; i++) {
        if (path->Buffer[i] == '~')
            return TRUE;
    }
    return FALSE;
}

BOOLEAN IsPathAllowed(
    _In_ PCUNICODE_STRING path)
{
    if (IsNullOrEmpty(path) || path->Length < PATH_MIN_SIZE_EXE) {
        return FALSE;
    }

    // check environment variable
    if (path->Buffer[1] != '?') {
        return TRUE;
    }

    // check C path:\??\C:*
    if (_wcsnicmp(path->Buffer + PATH_START_OFFSET, L"?\\C:", SIZE_OF_C) == 0) {

        //check [\??\C:\] Windows\*,ignore case
        if (_wcsnicmp(path->Buffer + PATH_FOLDER_OFFSET, L"Windows\\", SIZE_OF_WINDOWS) == 0) {
            return TRUE;
        }

		//check [\??\C:\] Progra~* or Program Files*,ignore case
        if (_wcsnicmp(path->Buffer + PATH_FOLDER_OFFSET, L"Progra", SIZE_OF_PROGRA) == 0) {
            if (path->Buffer[PATH_PROGRA_OFFSET] == '~')
                return TRUE;
            if (_wcsnicmp(path->Buffer + PATH_PROGRA_OFFSET, L"m Files", SIZE_OF_MFILES) == 0) {
                return TRUE;
            }
        }
    }
    return FALSE;
}